package com.milvus.imagesearch;

import io.milvus.v2.client.ConnectConfig;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.v2.service.vector.request.data.BaseVector;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.SearchResp;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * 图片搜索器
 */
@Service
public class ImageSearcher {

    private  MilvusClientV2 client=new MilvusClientV2(
            ConnectConfig.builder()
                    .uri("http://116.62.142.81:19530")
                    .token("root:Milvus")
                    .build()
    );

    public List<List<SearchResp.SearchResult>> search(String collectionName ,INDArray queryFeatures) {
        // 在Milvus中进行相似性搜索
        float[] floatArray = queryFeatures.toFloatVector();
        List<BaseVector> floatList = new ArrayList<>();

        FloatVec floatVec = new FloatVec(floatArray);
        floatList.add(floatVec);

        SearchReq searchReq = SearchReq.builder()
                .collectionName(collectionName)
                .data(floatList)
                .outputFields(Arrays.asList("name"))
                .topK(10)
                .offset(10)
                .build();

        SearchResp searchResp = client.search(searchReq);

        List<List<SearchResp.SearchResult>> searchResults = searchResp.getSearchResults();
        return searchResults;

    }
}
